#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
高胜率股票买卖点预测算法 - 增强版
只在高确定性信号时才提示买入/卖出
"""

import numpy as np
import pandas as pd
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass
from enum import Enum
import json

class Signal(Enum):
    """交易信号 - 只保留强信号"""
    STRONG_BUY = "强烈买入"
    STRONG_SELL = "强烈卖出"
    HOLD = "持有观望"

@dataclass
class TradingPoint:
    """交易点"""
    date: str
    price: float
    signal: Signal
    confidence: float  # 信心度 0-100
    reasons: List[str]  # 多个原因
    indicators: Dict[str, float]
    risk_level: str  # 风险等级

class EnhancedTechnicalAnalysis:
    """增强版技术分析"""
    
    @staticmethod
    def calculate_ma(prices: List[float], period: int) -> List[float]:
        """计算移动平均线"""
        if len(prices) < period:
            return []
        
        result = []
        for i in range(len(prices)):
            if i < period - 1:
                result.append(np.nan)
            else:
                ma = np.mean(prices[i-period+1:i+1])
                result.append(ma)
        return result
    
    @staticmethod
    def calculate_ema(prices: List[float], period: int) -> List[float]:
        """计算指数移动平均线"""
        if len(prices) < period:
            return []
        
        result = []
        multiplier = 2 / (period + 1)
        sma = np.mean(prices[:period])
        result.append(sma)
        
        for i in range(1, len(prices)):
            if i < period:
                result.append(np.nan)
            else:
                ema = (prices[i] - result[-1]) * multiplier + result[-1]
                result.append(ema)
        
        return result
    
    @staticmethod
    def calculate_macd(prices: List[float]) -> Dict[str, List[float]]:
        """计算MACD"""
        ema12 = EnhancedTechnicalAnalysis.calculate_ema(prices, 12)
        ema26 = EnhancedTechnicalAnalysis.calculate_ema(prices, 26)
        
        macd_line = [e12 - e26 if not np.isnan(e12) and not np.isnan(e26) else np.nan 
                     for e12, e26 in zip(ema12, ema26)]
        
        signal_line = EnhancedTechnicalAnalysis.calculate_ema(
            [x for x in macd_line if not np.isnan(x)], 9
        )
        signal_line = [np.nan] * (len(macd_line) - len(signal_line)) + signal_line
        
        histogram = [m - s if not np.isnan(m) and not np.isnan(s) else np.nan 
                     for m, s in zip(macd_line, signal_line)]
        
        return {'macd': macd_line, 'signal': signal_line, 'histogram': histogram}
    
    @staticmethod
    def calculate_rsi(prices: List[float], period: int = 14) -> List[float]:
        """计算RSI"""
        if len(prices) < period + 1:
            return []
        
        deltas = [prices[i] - prices[i-1] for i in range(1, len(prices))]
        gains = [d if d > 0 else 0 for d in deltas]
        losses = [-d if d < 0 else 0 for d in deltas]
        
        result = [np.nan]
        avg_gain = np.mean(gains[:period])
        avg_loss = np.mean(losses[:period])
        
        for i in range(period, len(deltas)):
            if avg_loss == 0:
                rsi = 100
            else:
                rs = avg_gain / avg_loss
                rsi = 100 - (100 / (1 + rs))
            
            result.append(rsi)
            avg_gain = (avg_gain * (period - 1) + gains[i]) / period
            avg_loss = (avg_loss * (period - 1) + losses[i]) / period
        
        return result
    
    @staticmethod
    def calculate_bollinger_bands(prices: List[float], period: int = 20) -> Dict[str, List[float]]:
        """计算布林带"""
        ma = EnhancedTechnicalAnalysis.calculate_ma(prices, period)
        upper, lower = [], []
        
        for i in range(len(prices)):
            if i < period - 1:
                upper.append(np.nan)
                lower.append(np.nan)
            else:
                std = np.std(prices[i-period+1:i+1])
                upper.append(ma[i] + 2 * std)
                lower.append(ma[i] - 2 * std)
        
        return {'upper': upper, 'middle': ma, 'lower': lower}
    
    @staticmethod
    def calculate_volume_ma(volumes: List[float], period: int = 5) -> List[float]:
        """计算成交量均线"""
        return EnhancedTechnicalAnalysis.calculate_ma(volumes, period)

class HighAccuracyPredictor:
    """高准确率预测器 - 只在强信号时才提示"""
    
    def __init__(self):
        self.ta = EnhancedTechnicalAnalysis()
        self.min_buy_score = 8  # 买入最低分数（提高门槛）
        self.min_sell_score = 8  # 卖出最低分数
        self.min_confidence = 70  # 最低信心度
    
    def analyze(self, stock_data: List[Dict]) -> List[TradingPoint]:
        """
        高准确率分析
        只返回强烈买入/强烈卖出信号
        """
        if len(stock_data) < 60:
            return []
        
        # 提取数据
        dates = [d['date'] for d in stock_data]
        opens = [float(d['open']) for d in stock_data]
        closes = [float(d['close']) for d in stock_data]
        highs = [float(d['high']) for d in stock_data]
        lows = [float(d['low']) for d in stock_data]
        volumes = [float(d['volume']) for d in stock_data]
        
        # 计算技术指标
        ma5 = self.ta.calculate_ma(closes, 5)
        ma10 = self.ta.calculate_ma(closes, 10)
        ma20 = self.ta.calculate_ma(closes, 20)
        ma60 = self.ta.calculate_ma(closes, 60)
        
        macd = self.ta.calculate_macd(closes)
        rsi = self.ta.calculate_rsi(closes, 14)
        bollinger = self.ta.calculate_bollinger_bands(closes, 20)
        volume_ma = self.ta.calculate_volume_ma(volumes, 5)
        
        trading_points = []
        
        # 从第60天开始分析（确保所有指标都有数据）
        for i in range(60, len(stock_data)):
            buy_score = 0
            sell_score = 0
            reasons = []
            
            # === 买入信号检测（严格条件） ===
            
            # 1. MACD金叉（强信号 +3分）
            if (i > 0 and 
                not np.isnan(macd['macd'][i]) and 
                not np.isnan(macd['signal'][i]) and
                macd['macd'][i-1] <= macd['signal'][i-1] and 
                macd['macd'][i] > macd['signal'][i] and
                macd['macd'][i] < 0):  # 在0轴下方金叉更可靠
                buy_score += 3
                reasons.append("MACD金叉（低位）")
            
            # 2. RSI超卖反弹（强信号 +3分）
            if i < len(rsi) and not np.isnan(rsi[i]):
                if rsi[i] < 30 and i > 0 and rsi[i] > rsi[i-1]:
                    buy_score += 3
                    reasons.append(f"RSI超卖反弹({rsi[i]:.1f})")
                elif rsi[i] < 25:
                    buy_score += 2
                    reasons.append(f"RSI深度超卖({rsi[i]:.1f})")
            
            # 3. 均线多头排列（+2分）
            if (not np.isnan(ma5[i]) and not np.isnan(ma10[i]) and 
                not np.isnan(ma20[i]) and not np.isnan(ma60[i])):
                if ma5[i] > ma10[i] > ma20[i] > ma60[i]:
                    buy_score += 2
                    reasons.append("完美多头排列")
                elif ma5[i] > ma10[i] > ma20[i]:
                    buy_score += 1
                    reasons.append("多头排列")
            
            # 4. 均线金叉（+2分）
            if (i > 0 and not np.isnan(ma5[i]) and not np.isnan(ma10[i]) and
                ma5[i-1] <= ma10[i-1] and ma5[i] > ma10[i]):
                buy_score += 2
                reasons.append("MA5金叉MA10")
            
            # 5. 布林带下轨反弹（+2分）
            if (not np.isnan(bollinger['lower'][i]) and 
                closes[i] < bollinger['lower'][i] * 1.01 and  # 接近或触及下轨
                i > 0 and closes[i] > closes[i-1]):  # 开始反弹
                buy_score += 2
                reasons.append("下轨反弹")
            
            # 6. 成交量温和放大（+1分）
            if i >= 5 and not np.isnan(volume_ma[i]):
                if volumes[i] > volume_ma[i] * 1.2 and closes[i] > opens[i]:
                    buy_score += 1
                    reasons.append("温和放量上涨")
            
            # 7. 价格位置（+1分）
            if not np.isnan(ma60[i]) and closes[i] < ma60[i] * 0.95:
                buy_score += 1
                reasons.append("价格低位")
            
            # 8. 连续下跌后反弹（+1分）
            if i >= 3:
                if all(closes[i-j] < closes[i-j-1] for j in range(3)) and closes[i] > closes[i-1]:
                    buy_score += 1
                    reasons.append("连续下跌后反弹")
            
            # === 卖出信号检测（严格条件） ===
            
            # 1. MACD死叉（强信号 +3分）
            if (i > 0 and 
                not np.isnan(macd['macd'][i]) and 
                not np.isnan(macd['signal'][i]) and
                macd['macd'][i-1] >= macd['signal'][i-1] and 
                macd['macd'][i] < macd['signal'][i] and
                macd['macd'][i] > 0):  # 在0轴上方死叉
                sell_score += 3
                reasons.append("MACD死叉（高位）")
            
            # 2. RSI超买回落（强信号 +3分）
            if i < len(rsi) and not np.isnan(rsi[i]):
                if rsi[i] > 70 and i > 0 and rsi[i] < rsi[i-1]:
                    sell_score += 3
                    reasons.append(f"RSI超买回落({rsi[i]:.1f})")
                elif rsi[i] > 75:
                    sell_score += 2
                    reasons.append(f"RSI严重超买({rsi[i]:.1f})")
            
            # 3. 均线空头排列（+2分）
            if (not np.isnan(ma5[i]) and not np.isnan(ma10[i]) and 
                not np.isnan(ma20[i]) and not np.isnan(ma60[i])):
                if ma5[i] < ma10[i] < ma20[i] < ma60[i]:
                    sell_score += 2
                    reasons.append("完美空头排列")
                elif ma5[i] < ma10[i] < ma20[i]:
                    sell_score += 1
                    reasons.append("空头排列")
            
            # 4. 均线死叉（+2分）
            if (i > 0 and not np.isnan(ma5[i]) and not np.isnan(ma10[i]) and
                ma5[i-1] >= ma10[i-1] and ma5[i] < ma10[i]):
                sell_score += 2
                reasons.append("MA5死叉MA10")
            
            # 5. 布林带上轨回落（+2分）
            if (not np.isnan(bollinger['upper'][i]) and 
                closes[i] > bollinger['upper'][i] * 0.99 and  # 接近或突破上轨
                i > 0 and closes[i] < closes[i-1]):  # 开始回落
                sell_score += 2
                reasons.append("上轨回落")
            
            # 6. 高位放量滞涨（+2分）
            if (i >= 5 and not np.isnan(volume_ma[i]) and not np.isnan(ma20[i])):
                if (volumes[i] > volume_ma[i] * 1.5 and 
                    closes[i] > ma20[i] * 1.1 and
                    abs(closes[i] - opens[i]) / opens[i] < 0.01):  # 涨幅很小
                    sell_score += 2
                    reasons.append("高位放量滞涨")
            
            # 7. 连续上涨后回落（+1分）
            if i >= 3:
                if all(closes[i-j] > closes[i-j-1] for j in range(3)) and closes[i] < closes[i-1]:
                    sell_score += 1
                    reasons.append("连续上涨后回落")
            
            # === 生成信号（只有强信号） ===
            
            # 强烈买入：分数>=8，且无卖出信号
            if buy_score >= self.min_buy_score and sell_score == 0:
                confidence = min(buy_score * 10, 95)
                
                if confidence >= self.min_confidence:
                    # 评估风险
                    risk_level = "低风险" if buy_score >= 10 else "中等风险"
                    
                    trading_point = TradingPoint(
                        date=dates[i],
                        price=closes[i],
                        signal=Signal.STRONG_BUY,
                        confidence=confidence,
                        reasons=reasons[:4],  # 最多4个原因
                        risk_level=risk_level,
                        indicators={
                            'ma5': ma5[i] if not np.isnan(ma5[i]) else 0,
                            'ma10': ma10[i] if not np.isnan(ma10[i]) else 0,
                            'ma20': ma20[i] if not np.isnan(ma20[i]) else 0,
                            'rsi': rsi[i] if i < len(rsi) and not np.isnan(rsi[i]) else 0,
                            'macd': macd['macd'][i] if not np.isnan(macd['macd'][i]) else 0,
                            'buy_score': buy_score,
                            'sell_score': sell_score
                        }
                    )
                    trading_points.append(trading_point)
            
            # 强烈卖出：分数>=8，且无买入信号
            elif sell_score >= self.min_sell_score and buy_score == 0:
                confidence = min(sell_score * 10, 95)
                
                if confidence >= self.min_confidence:
                    risk_level = "高风险" if sell_score >= 10 else "中等风险"
                    
                    trading_point = TradingPoint(
                        date=dates[i],
                        price=closes[i],
                        signal=Signal.STRONG_SELL,
                        confidence=confidence,
                        reasons=reasons[:4],
                        risk_level=risk_level,
                        indicators={
                            'ma5': ma5[i] if not np.isnan(ma5[i]) else 0,
                            'ma10': ma10[i] if not np.isnan(ma10[i]) else 0,
                            'ma20': ma20[i] if not np.isnan(ma20[i]) else 0,
                            'rsi': rsi[i] if i < len(rsi) and not np.isnan(rsi[i]) else 0,
                            'macd': macd['macd'][i] if not np.isnan(macd['macd'][i]) else 0,
                            'buy_score': buy_score,
                            'sell_score': sell_score
                        }
                    )
                    trading_points.append(trading_point)
        
        return trading_points
    
    def get_current_recommendation(self, stock_data: List[Dict]) -> Dict:
        """获取当前推荐"""
        if len(stock_data) < 60:
            return {
                'signal': 'HOLD',
                'signal_cn': '持有观望',
                'confidence': 0,
                'reasons': ['数据不足'],
                'risk_level': '未知',
                'price': stock_data[-1]['close'] if stock_data else 0
            }
        
        trading_points = self.analyze(stock_data)
        
        if not trading_points:
            return {
                'signal': 'HOLD',
                'signal_cn': '持有观望',
                'confidence': 50,
                'reasons': ['无强烈信号，建议观望'],
                'risk_level': '中等',
                'price': stock_data[-1]['close']
            }
        
        # 返回最新的强信号
        latest = trading_points[-1]
        return {
            'signal': latest.signal.name,
            'signal_cn': latest.signal.value,
            'confidence': latest.confidence,
            'reasons': latest.reasons,
            'risk_level': latest.risk_level,
            'price': latest.price,
            'date': latest.date,
            'indicators': latest.indicators
        }
    
    def backtest(self, stock_data: List[Dict], initial_capital: float = 100000) -> Dict:
        """回测 - 只使用强信号"""
        trading_points = self.analyze(stock_data)
        
        capital = initial_capital
        shares = 0
        trades = []
        max_capital = initial_capital
        max_drawdown = 0
        
        for point in trading_points:
            current_value = capital + shares * point.price
            
            if point.signal == Signal.STRONG_BUY and capital > 0:
                # 强烈买入 - 使用80%资金
                invest_amount = capital * 0.8
                shares_to_buy = int(invest_amount / point.price)
                
                if shares_to_buy > 0:
                    cost = shares_to_buy * point.price
                    capital -= cost
                    shares += shares_to_buy
                    trades.append({
                        'date': point.date,
                        'action': 'BUY',
                        'price': point.price,
                        'shares': shares_to_buy,
                        'confidence': point.confidence,
                        'reasons': point.reasons
                    })
            
            elif point.signal == Signal.STRONG_SELL and shares > 0:
                # 强烈卖出 - 全部卖出
                revenue = shares * point.price
                capital += revenue
                trades.append({
                    'date': point.date,
                    'action': 'SELL',
                    'price': point.price,
                    'shares': shares,
                    'confidence': point.confidence,
                    'reasons': point.reasons
                })
                shares = 0
            
            # 更新最大资金和最大回撤
            current_value = capital + shares * point.price
            max_capital = max(max_capital, current_value)
            drawdown = (max_capital - current_value) / max_capital
            max_drawdown = max(max_drawdown, drawdown)
        
        # 计算最终价值
        final_price = stock_data[-1]['close']
        final_value = capital + shares * final_price
        profit = final_value - initial_capital
        profit_rate = (profit / initial_capital) * 100
        
        # 买入持有策略
        buy_hold_shares = int(initial_capital / stock_data[0]['close'])
        buy_hold_value = buy_hold_shares * final_price
        buy_hold_profit_rate = ((buy_hold_value - initial_capital) / initial_capital) * 100
        
        # 计算胜率
        win_count = 0
        total_profit_trades = 0
        
        for i in range(1, len(trades)):
            if trades[i]['action'] == 'SELL' and trades[i-1]['action'] == 'BUY':
                if trades[i]['price'] > trades[i-1]['price']:
                    win_count += 1
                total_profit_trades += 1
        
        win_rate = (win_count / total_profit_trades * 100) if total_profit_trades > 0 else 0
        
        return {
            'initial_capital': initial_capital,
            'final_value': final_value,
            'profit': profit,
            'profit_rate': profit_rate,
            'trades_count': len(trades),
            'trades': trades,
            'buy_hold_rate': buy_hold_profit_rate,
            'excess_return': profit_rate - buy_hold_profit_rate,
            'win_rate': win_rate,
            'max_drawdown': max_drawdown * 100,
            'total_profit_trades': total_profit_trades
        }

def main():
    """测试高准确率预测算法"""
    print("="*80)
    print("📊 高胜率股票预测算法 - 增强版测试")
    print("="*80)
    print("\n策略说明:")
    print("  • 只在强烈买入时才提示买入（买入分数≥8分）")
    print("  • 只在强烈卖出时才提示卖出（卖出分数≥8分）")
    print("  • 信心度必须≥70%")
    print("  • 严格的多指标验证")
    print("="*80)
    
    # 生成测试数据
    np.random.seed(42)
    dates = pd.date_range(start='2024-01-01', periods=120, freq='D')
    
    base_price = 100
    prices = [base_price]
    for i in range(119):
        change = np.random.randn() * 1.5
        new_price = prices[-1] + change
        prices.append(max(new_price, 1))
    
    stock_data = []
    for i, date in enumerate(dates):
        stock_data.append({
            'date': date.strftime('%Y-%m-%d'),
            'open': prices[i] * (1 + np.random.uniform(-0.01, 0.01)),
            'close': prices[i],
            'high': prices[i] * (1 + abs(np.random.uniform(0, 0.02))),
            'low': prices[i] * (1 - abs(np.random.uniform(0, 0.02))),
            'volume': np.random.randint(10000, 100000)
        })
    
    # 创建预测器
    predictor = HighAccuracyPredictor()
    
    # 分析
    print("\n📈 分析交易信号...")
    trading_points = predictor.analyze(stock_data)
    
    print(f"\n找到 {len(trading_points)} 个强信号:\n")
    if trading_points:
        print(f"{'日期':<12} {'价格':<10} {'信号':<12} {'信心度':<8} {'风险':<10} {'原因':<50}")
        print("-" * 105)
        
        for point in trading_points:
            reasons_str = " | ".join(point.reasons)
            print(f"{point.date:<12} {point.price:<10.2f} {point.signal.value:<12} "
                  f"{point.confidence:<8.1f} {point.risk_level:<10} {reasons_str:<50}")
    
    # 当前建议
    print("\n" + "="*80)
    print("💡 当前交易建议")
    print("="*80)
    
    recommendation = predictor.get_current_recommendation(stock_data)
    print(f"\n信号: {recommendation['signal_cn']}")
    print(f"信心度: {recommendation['confidence']:.1f}%")
    print(f"风险等级: {recommendation['risk_level']}")
    print(f"价格: ¥{recommendation['price']:.2f}")
    print(f"原因:")
    for reason in recommendation['reasons']:
        print(f"  • {reason}")
    
    # 回测
    print("\n" + "="*80)
    print("🔬 策略回测（高胜率模式）")
    print("="*80)
    
    backtest_result = predictor.backtest(stock_data, 100000)
    
    print(f"\n📊 回测结果:")
    print(f"  初始资金: ¥{backtest_result['initial_capital']:,.0f}")
    print(f"  最终价值: ¥{backtest_result['final_value']:,.0f}")
    print(f"  盈亏金额: ¥{backtest_result['profit']:,.0f}")
    print(f"  收益率: {backtest_result['profit_rate']:.2f}%")
    print(f"  交易次数: {backtest_result['trades_count']}")
    print(f"  胜率: {backtest_result['win_rate']:.1f}%")
    print(f"  最大回撤: {backtest_result['max_drawdown']:.2f}%")
    
    print(f"\n📈 对比基准:")
    print(f"  买入持有收益率: {backtest_result['buy_hold_rate']:.2f}%")
    print(f"  超额收益: {backtest_result['excess_return']:.2f}%")
    
    if backtest_result['trades']:
        print(f"\n📝 交易记录:")
        print(f"{'日期':<12} {'操作':<6} {'价格':<10} {'数量':<8} {'信心度':<8} {'原因':<40}")
        print("-" * 90)
        for trade in backtest_result['trades']:
            reasons_str = " | ".join(trade['reasons'][:2])
            print(f"{trade['date']:<12} {trade['action']:<6} {trade['price']:<10.2f} "
                  f"{trade['shares']:<8} {trade['confidence']:<8.1f} {reasons_str:<40}")
    
    # 导出结果
    output = {
        'strategy': 'high_accuracy',
        'min_buy_score': predictor.min_buy_score,
        'min_confidence': predictor.min_confidence,
        'trading_points': [
            {
                'date': p.date,
                'price': p.price,
                'signal': p.signal.value,
                'confidence': p.confidence,
                'reasons': p.reasons,
                'risk_level': p.risk_level,
                'indicators': p.indicators
            }
            for p in trading_points
        ],
        'current_recommendation': recommendation,
        'backtest': backtest_result
    }
    
    with open('high_accuracy_prediction.json', 'w', encoding='utf-8') as f:
        json.dump(output, f, ensure_ascii=False, indent=2)
    
    print(f"\n💾 结果已保存到 high_accuracy_prediction.json")
    print("\n" + "="*80)
    print("✅ 测试完成！")
    print("="*80)

if __name__ == "__main__":
    main()


